#include <iostream>
#include <opencv2/opencv.hpp>
#include "ED_Lib/EDLib.h"
#include <ctime>

using namespace cv;
using namespace std;

int main()
{
    cv::Mat cameraMatrix;
    cameraMatrix = cv::Mat::zeros(3, 3, CV_64F);
    std::vector<double> distcoeffs{0.136717898475962, -0.253882847667133, 0.0, 0.0, 0.0};
    cameraMatrix.at<double>(0, 0) = 893.570013573463;
    cameraMatrix.at<double>(0, 2) = 632.175640734913;
    cameraMatrix.at<double>(1, 1) = 894.547565226624;
    cameraMatrix.at<double>(1, 2) = 493.266669326202;
    cameraMatrix.at<double>(2, 2) = 1;

    cv::VideoCapture cap("/home/ljy/Videos/dji_red_circle_video/DJI_0003.MP4");
    while(1)
    {
        clock_t start_time, end_time;
        start_time = clock();
        cv::Mat frame;
        cap >> frame;
        cv::resize(frame, frame, cv::Size(frame.size().width / 2, frame.size().height / 2));
        imshow("frame", frame);
        cv::cvtColor(frame,frame, COLOR_BGR2YUV);
        std::vector<cv::Mat> channels;
        cv::split(frame, channels);
        cv::Mat testImage(frame.size(), CV_8UC1);
        testImage = channels[2];

        EDCircles testEDCircles = EDCircles(testImage);
        Mat circleImg = testEDCircles.drawResult(true, ImageStyle::BOTH);
//        //Get circle information as [cx, cy, r]
        vector<mCircle> circles = testEDCircles.getCircles();
        mCircle* p_biggest_circle = NULL;
        if(circles.size() >= 1)
        {
            p_biggest_circle = &circles[0];
            for(int i = 1; i < circles.size(); i++)
            {
                if(circles[i].r > p_biggest_circle->r)
                    p_biggest_circle = &circles[i];
            }
        }
        vector<mEllipse> ellipses = testEDCircles.getEllipses();
        mEllipse* p_biggest_ellipse = NULL;
        if(ellipses.size() >= 1)
        {
            p_biggest_ellipse = &ellipses[0];
            for(int i = 1; i < ellipses.size(); i++)
            {
                if(ellipses[i].axes.area() > p_biggest_ellipse->axes.area())
                    p_biggest_ellipse = &ellipses[i];
            }
        }

        vector<cv::Point2d> four_points; // up down left right
        if( p_biggest_circle == NULL && p_biggest_ellipse == NULL)
            cout << " No Circle!!" << endl;
        else if( p_biggest_ellipse == NULL ||
                    (p_biggest_circle != NULL && p_biggest_ellipse != NULL && p_biggest_circle->r * p_biggest_circle->r > p_biggest_ellipse->axes.area()))
        {
            four_points.emplace_back(p_biggest_circle->center.x, p_biggest_circle->center.y - p_biggest_circle->r);
            four_points.emplace_back(p_biggest_circle->center.x, p_biggest_circle->center.y + p_biggest_circle->r);
            four_points.emplace_back(p_biggest_circle->center.x - p_biggest_circle->r, p_biggest_circle->center.y);
            four_points.emplace_back(p_biggest_circle->center.x + p_biggest_circle->r, p_biggest_circle->center.y);
        } else if( p_biggest_circle == NULL ||
                   (p_biggest_circle != NULL && p_biggest_ellipse != NULL && p_biggest_circle->r * p_biggest_circle->r < p_biggest_ellipse->axes.area()))
        {
            four_points.emplace_back(p_biggest_ellipse->center.x,
                                     p_biggest_ellipse->center.y - sqrt(pow(p_biggest_ellipse->axes.width * cos(PI/2 - p_biggest_ellipse->theta), 2) +
                                                                        pow(p_biggest_ellipse->axes.height * sin(PI/2 - p_biggest_ellipse->theta), 2)));
            four_points.emplace_back(p_biggest_ellipse->center.x,
                                     p_biggest_ellipse->center.y + sqrt(pow(p_biggest_ellipse->axes.width * cos(PI/2 - p_biggest_ellipse->theta), 2) +
                                                                        pow(p_biggest_ellipse->axes.height * sin(PI/2 - p_biggest_ellipse->theta), 2)));
            four_points.emplace_back(p_biggest_ellipse->center.x - sqrt(pow(p_biggest_ellipse->axes.width * cos(PI - p_biggest_ellipse->theta), 2) +
                                                                        pow(p_biggest_ellipse->axes.height * sin(PI - p_biggest_ellipse->theta), 2)),
                                     p_biggest_ellipse->center.y );
            four_points.emplace_back(p_biggest_ellipse->center.x + sqrt(pow(p_biggest_ellipse->axes.width * cos(PI - p_biggest_ellipse->theta), 2) +
                                                                        pow(p_biggest_ellipse->axes.height * sin(PI - p_biggest_ellipse->theta), 2)),
                                     p_biggest_ellipse->center.y );
        }


        for(int i = 0; i < four_points.size(); i++)
        {
            circle(circleImg, four_points[i], 2, cv::Scalar(0,255,255),2);
            putText(circleImg, to_string(i), four_points[i], cv::FONT_HERSHEY_PLAIN, 2, cv::Scalar(0,255,255));

        }

        imshow("Circle Image 1", circleImg);

        std::vector<cv::Point3d> big_circle_object_points({cv::Point3d(0,0,750),
                                                           cv::Point3d(0,0,-750),
                                                           cv::Point3d(0,750,0),
                                                           cv::Point3d(0,-750,0)});
        std::vector<cv::Point3d> small_circle_object_points({cv::Point3d(0,0,600),
                                                           cv::Point3d(0,0,-600),
                                                           cv::Point3d(0,600,0),
                                                           cv::Point3d(0,-600,0)});
        cv::Mat big_rvec, big_tvec, small_rvec, small_tvec;
        cv::solvePnP(big_circle_object_points, four_points, cameraMatrix, distcoeffs, big_rvec, big_tvec);
        cv::solvePnP(small_circle_object_points, four_points, cameraMatrix, distcoeffs, small_rvec, small_tvec);

        end_time = clock();
//        cout << "Time: " << double(end_time - start_time) / CLOCKS_PER_SEC << endl;
        cout << big_tvec << endl;
        char c = waitKey(1);
        if(c == 27) break;
    }
    return 0;
}



